package com.schoolai.util.tools;
import com.nimbusds.jose.*;
import com.nimbusds.jose.crypto.RSASSASigner;
import com.nimbusds.jose.crypto.RSASSAVerifier;
import com.nimbusds.jwt.JWTClaimsSet;
import com.nimbusds.jwt.SignedJWT;

import java.io.*;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.security.*;
import java.security.cert.Certificate;
import java.security.cert.CertificateException;
import java.security.interfaces.RSAPrivateKey;
import java.security.interfaces.RSAPublicKey;
import java.text.ParseException;
import java.util.Base64;
import java.util.Date;
import java.util.UUID;
/**
 * RSA非对称加密JWT工具类
 * 从JKS文件加载私钥，从文本文件加载公钥
 */

public class JWTUtil {
    /**
     * 从JKS文件加载私钥
     * @param fis JKS文件
     * @param alias 密钥别名
     * @param password 密钥密码
     * @return RSAPrivateKey对象
     */
    public static RSAPrivateKey loadPrivateKeyFromJKS(FileInputStream fis, String alias, String password)
            throws KeyStoreException, IOException, NoSuchAlgorithmException, CertificateException, UnrecoverableKeyException {

        KeyStore keyStore = KeyStore.getInstance("JKS");
        keyStore.load(fis, password.toCharArray());
        Key key = keyStore.getKey(alias, password.toCharArray());
        if (key instanceof RSAPrivateKey) {
            return (RSAPrivateKey) key;
        } else {
            throw new IllegalArgumentException("Key is not an RSA private key");
        }
    }

    /**
     * 从文本文件加载公钥
     * @param publicKeyKeyContent 公钥文件路径
     * @return RSAPublicKey对象
     */
    public static RSAPublicKey loadPublicKeyFromFile(String publicKeyKeyContent)
            throws IOException, GeneralSecurityException {

        //String keyContent = new String(Files.readAllBytes(Paths.get(publicKeyFilePath)));
        String keyContent =  publicKeyKeyContent;
        keyContent = keyContent.replace("-----BEGIN PUBLIC KEY-----", "")
                .replace("-----END PUBLIC KEY-----", "")
                .replaceAll("\\s", "");

        byte[] keyBytes = Base64.getDecoder().decode(keyContent);

        KeyFactory keyFactory = KeyFactory.getInstance("RSA");
        return (RSAPublicKey) keyFactory.generatePublic(new java.security.spec.X509EncodedKeySpec(keyBytes));
    }

    /**
     * 从JKS文件加载公钥（通过证书）
     * @param jksFilePath JKS文件路径
     * @param alias 密钥别名
     * @param password 密钥密码
     * @return RSAPublicKey对象
     */
    public static RSAPublicKey loadPublicKeyFromJKS(String jksFilePath, String alias, String password)
            throws KeyStoreException, IOException, NoSuchAlgorithmException, CertificateException {

        KeyStore keyStore = KeyStore.getInstance("JKS");
        try (FileInputStream fis = new FileInputStream(jksFilePath)) {
            keyStore.load(fis, password.toCharArray());
        }

        Certificate cert = keyStore.getCertificate(alias);
        if (cert != null) {
            PublicKey publicKey = cert.getPublicKey();
            if (publicKey instanceof RSAPublicKey) {
                return (RSAPublicKey) publicKey;
            }
        }

        throw new IllegalArgumentException("No RSA public key found for alias: " + alias);
    }

    /**
     * 生成JWT令牌
     * @param subject 主题（通常是用户ID）
     * @param issuer 签发者
     * @param expirationTime 过期时间（毫秒）
     * @param claims 自定义声明
     * @return JWT令牌字符串
     * @throws JOSEException
     */
    public String generateToken(RSAPrivateKey rsaPrivateKey,String subject, String issuer, long expirationTime,
                                java.util.Map<String, Object> claims) throws JOSEException {
        // 创建JWT头部，指定RS256算法
        JWSHeader header = new JWSHeader.Builder(JWSAlgorithm.RS256)
                .type(JOSEObjectType.JWT)
                .keyID(UUID.randomUUID().toString())
                .build();

        // 构建JWT声明集
        JWTClaimsSet.Builder claimsBuilder = new JWTClaimsSet.Builder()
                .subject(subject)
                .issuer(issuer)
                .issueTime(new Date())
                .expirationTime(new Date(System.currentTimeMillis() + expirationTime))
                .jwtID(UUID.randomUUID().toString());

        // 添加自定义声明
        if (claims != null) {
            for (java.util.Map.Entry<String, Object> entry : claims.entrySet()) {
                claimsBuilder.claim(entry.getKey(), entry.getValue());
            }
        }

        JWTClaimsSet claimsSet = claimsBuilder.build();

        // 创建签名器
        JWSSigner signer = new RSASSASigner(rsaPrivateKey);

        // 创建签名JWT
        SignedJWT signedJWT = new SignedJWT(header, claimsSet);

        // 应用签名
        signedJWT.sign(signer);

        // 序列化为字符串
        return signedJWT.serialize();
    }

    /**
     * 验证JWT令牌
     * @param token JWT令牌字符串
     * @return 验证结果和声明集
     * @throws ParseException
     * @throws JOSEException
     */
    public VerificationResult verifyToken(RSAPublicKey rsaPublicKey,String token) throws ParseException, JOSEException {
        // 解析JWT
        SignedJWT signedJWT = SignedJWT.parse(token);

        // 创建验证器
        JWSVerifier verifier = new RSASSAVerifier(rsaPublicKey);

        // 验证签名
        boolean signatureValid = signedJWT.verify(verifier);

        if (!signatureValid) {
            return new VerificationResult(false, null, "Invalid signature");
        }

        // 验证过期时间
        JWTClaimsSet claimsSet = signedJWT.getJWTClaimsSet();
        Date expirationTime = claimsSet.getExpirationTime();

        if (expirationTime != null && expirationTime.before(new Date())) {
            return new VerificationResult(false, null, "Token expired");
        }

        return new VerificationResult(true, claimsSet, "Token is valid");
    }

    /**
     * 验证结果封装类
     */
    public static class VerificationResult {
        private final boolean valid;
        private final JWTClaimsSet claims;
        private final String message;

        public VerificationResult(boolean valid, JWTClaimsSet claims, String message) {
            this.valid = valid;
            this.claims = claims;
            this.message = message;
        }

        public boolean isValid() {
            return valid;
        }

        public JWTClaimsSet getClaims() {
            return claims;
        }

        public String getMessage() {
            return message;
        }
    }
}
